import numpy as np

TRAJ_LEN = 200
G = 6.67430e-11  # m^3 kg^-1 s^-2
STEP = 3600

class Planet:
    
    def __init__(self, name, mass, pos, vel):
        self.name = name
        self.mass = mass
        self.forces = np.zeros(3, dtype='float64')
        self.pos = pos
        # print(self.pos)
        self.vel = vel
        self.trajectory = np.array([pos])

    
    def update(self, acc):
        self.vel += acc * STEP
        self.pos += self.vel * STEP
        new_traj = np.concatenate([self.trajectory, [self.pos]]) 
        self.trajectory = new_traj[-1*TRAJ_LEN:]
        # self.trajectory = np.concatenate([self.trajectory, [self.pos]]) 


class PlanetManager:
    def __init__(self):
        
        self.planets = [
            Planet('A', 5.972e24, np.array([0, 0, 0], dtype='float64'), np.array([1, 0, 0], dtype='float64')),
            Planet('B', 7.348e23, np.array([3.8e8, 0, 0], dtype='float64'), np.array([10, 1022, 0], dtype='float64')),
            Planet('C', 7.348e23, np.array([0, 2.57e8, 3.8e8], dtype='float64'), np.array([50, 70, 1020], dtype='float64')),
            # Planet('D', 6.4171e23, np.array([0, 2.02e8, 1.07e7], dtype='float64'), np.array([1030, 20, 0], dtype='float64')),
            ]
        
    def calcute_net_force(self, planet):

        for other in self.planets:
            if other.name == planet.name:
                continue

            r = other.pos - planet.pos
            # print(planet.name, planet.pos)

            r_len = np.linalg.norm(r, 2)

            F = G * planet.mass * other.mass * r/ r_len ** 3
            # print(planet.name, other.name, F)

            acc1 = F / planet.mass
            acc2 = -F / other.mass

            planet.forces += acc1
            other.forces += acc2


        
    def update_state(self):
        '''
            完成一次运算
        '''

        # 初始化所有星体力表
        for planet in self.planets:
            planet.forces = np.zeros(3, dtype='float64')
            # print('for planet in self.planets:',planet.name, planet.pos)
        
        # 遍历星体，分别计算该星体受其他星体的牵引力，以及该星体对其他星体的牵引力
        for planet in self.planets:
            # print(planet.name, planet.pos)
            self.calcute_net_force(planet)
        
        # 遍历星体，计算合力，并更新
        for planet in self.planets:
            # print(planet.name, planet.forces)
            # acc = planet.forces #/ planet.mass
            # print(planet.name, acc)
            planet.update(np.copy(planet.forces))

        
        

